#include "VulkanApplication.h"
#include <vector>
#include <iostream>
#include <fstream>
#include <cstring>

struct Vertex
{
    float x, y, z;
    float r, g, b;
};

bool createBuffer(VulkanApplication& app, VkBuffer& buffer, VkDeviceMemory& bufferMemory,
                  size_t dataSize, VkBufferUsageFlags usage, uint32_t propertyBits)
{
    VkBufferCreateInfo bufferInfo
    {
        VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
        nullptr, 0,
        dataSize,
        usage,
        VK_SHARING_MODE_EXCLUSIVE,
        0, nullptr  // pQueueFamilyIndices
    };

    VkResult result = vkCreateBuffer(app.getDevice(), &bufferInfo, nullptr, &buffer);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create data buffer.\n");
        return false;
    }

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(app.getDevice(), buffer, &memRequirements);

    VkPhysicalDeviceMemoryProperties memProperties;
    vkGetPhysicalDeviceMemoryProperties(app.getPhysicalDevice(), &memProperties);

    uint32_t memoryTypeIndex = 0;
    for (uint32_t i = 0; i < memProperties.memoryTypeCount; ++i)
    {
        if ((memRequirements.memoryTypeBits & (1 << i)) &&
            (memProperties.memoryTypes[i].propertyFlags & propertyBits) == propertyBits)
        { memoryTypeIndex = i; }
    }

    VkMemoryAllocateInfo allocInfo
    {
        VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
        nullptr,
        memRequirements.size,
        memoryTypeIndex
    };

    result = vkAllocateMemory(app.getDevice(), &allocInfo, nullptr, &bufferMemory);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create vertex buffer memory.\n");
        return false;
    }

    vkBindBufferMemory(app.getDevice(), buffer, bufferMemory, 0);
    return true;
}

VkBuffer createVertexBuffer(VulkanApplication& app, const std::vector<Vertex>& vertices)
{
    size_t dataSize = sizeof(Vertex) * vertices.size();
    VkBuffer vertexBuffer = nullptr;
    VkDeviceMemory vertexBufferMemory = nullptr;
    if (!createBuffer(app, vertexBuffer, vertexBufferMemory, dataSize, VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,
                      VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))
    { return nullptr; }

    void* data = nullptr;
    vkMapMemory(app.getDevice(), vertexBufferMemory, 0, dataSize, 0, &data);
    memcpy(data, vertices.data(), dataSize);
    vkUnmapMemory(app.getDevice(), vertexBufferMemory);
    return vertexBuffer;
}

VkBuffer createIndexBuffer(VulkanApplication& app, const std::vector<unsigned int>& indices)
{
    size_t dataSize = sizeof(unsigned int) * indices.size();
    VkBuffer indexBuffer = nullptr;
    VkDeviceMemory indexBufferMemory = nullptr;
    if (!createBuffer(app, indexBuffer, indexBufferMemory, dataSize, VK_BUFFER_USAGE_INDEX_BUFFER_BIT,
                      VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))
    { return nullptr; }

    void* data = nullptr;
    vkMapMemory(app.getDevice(), indexBufferMemory, 0, dataSize, 0, &data);
    memcpy(data, indices.data(), dataSize);
    vkUnmapMemory(app.getDevice(), indexBufferMemory);
    return indexBuffer;
}

int main(int argc, char** argv)
{
    VulkanApplication app("c11_indices", 1024, 768);
    if (!app.initialized())
    {
        printf("Failed to initialize VulkanApplication.\n");
        return 1;
    }

    // Create vertex & assembly data
    VkVertexInputBindingDescription bindingDescription
    {
        0, sizeof(Vertex),
        VK_VERTEX_INPUT_RATE_VERTEX
    };

    std::vector<VkVertexInputAttributeDescription> attrDescription;
    attrDescription.push_back(
        {
            0, 0,  // location, binding
            VK_FORMAT_R32G32B32_SFLOAT,  // format
            0                         // offset
        });
    attrDescription.push_back(
        {
            1, 0,  // location, binding
            VK_FORMAT_R32G32B32_SFLOAT,  // format
            sizeof(float) * 3         // offset
        });

    VkPipelineVertexInputStateCreateInfo inputVertexInfo
    {
        VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
        nullptr, 0,
        1, &bindingDescription,
        (uint32_t)attrDescription.size(), attrDescription.data()
    };

    VkPipelineInputAssemblyStateCreateInfo inputAssemblyInfo
    {
        VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
        nullptr, 0,
        VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST, VK_FALSE
    };

    // Create graphics pipeline
    if (!app.setUpPipeline("draw_vertices.vert", "draw_vertices.frag",
                           inputVertexInfo, inputAssemblyInfo))
    {
        printf("Failed to create pipeline.\n");
        return 1;
    }

    // Create vertex & index buffer
    std::vector<Vertex> vertices;
    vertices.push_back({ -0.5f, -0.5f, 0.0f, 0.0f, 1.0f, 0.0f });
    vertices.push_back({ 0.5f, -0.5f, 0.0f, 0.0f, 0.0f, 1.0f });
    vertices.push_back({ 0.5f, 0.5f, 0.0f, 1.0f, 0.0f, 0.0f });
    vertices.push_back({ -0.5f, 0.5f, 0.0f, 1.0f, 1.0f, 0.0f });

    std::vector<unsigned int> indices;
    indices.push_back(0); indices.push_back(1); indices.push_back(2);
    indices.push_back(0); indices.push_back(2); indices.push_back(3);

    VkBuffer vertexBuffer = createVertexBuffer(app, vertices);
    VkBuffer indexBuffer = createIndexBuffer(app, indices);

    // Main loop
    int width = 1024, height = 768;
    VkFence waitFence = app.getOrCreateFence("CommandFence");
    VkSemaphore waitNextImage = app.getOrCreateSemaphore("WaitNextImage");
    VkSemaphore waitSubmission = app.getOrCreateSemaphore("WaitSubmission");
    VkCommandBuffer commandBuffer = app.getCommandBuffer();
    while (!app.shouldClose())
    {
        vkWaitForFences(app.getDevice(), 1, &waitFence, VK_TRUE, UINT64_MAX);
        vkResetFences(app.getDevice(), 1, &waitFence);

        uint32_t imageIndex = 0;
        VkResult result = vkAcquireNextImageKHR(app.getDevice(), app.getSwapchain(), UINT64_MAX,
                                                waitNextImage, VK_NULL_HANDLE, &imageIndex);
        switch (result)
        {
        case VK_SUBOPTIMAL_KHR:
        case VK_ERROR_OUT_OF_DATE_KHR:
            imageIndex = app.recreateSwapchain(width, height, waitNextImage); break;
        case VK_SUCCESS:
            break;
        default:
            printf("Next image: %d (RESULT = %d)\n", imageIndex, result);
            app.setClosed(); break;
        }

        VkCommandBufferBeginInfo beginInfo
        {
            VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
            nullptr, 0, nullptr
        };
        vkResetCommandBuffer(commandBuffer, 0);
        vkBeginCommandBuffer(commandBuffer, &beginInfo);
        {
            VkClearValue clearColor = { {0.2f, 0.2f, 0.4f, 1.0f} };
            VkRenderPassBeginInfo renderPassInfo
            {
                VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO,
                nullptr, app.getRenderPass(),
                app.getSwapchainFramebuffers()[imageIndex],
                { {0, 0}, {(uint32_t)width, (uint32_t)height} },
                1, &clearColor
            };
            vkCmdBeginRenderPass(commandBuffer, &renderPassInfo, VK_SUBPASS_CONTENTS_INLINE);
            {
                VkViewport viewport
                { 0.0f, 0.0f, (float)width, (float)height, 0.0f, 1.0f };

                VkRect2D scissor
                { {0, 0}, {(uint32_t)width, (uint32_t)height} };

                vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, app.getPipeline());
                vkCmdSetViewport(commandBuffer, 0, 1, &viewport);
                vkCmdSetScissor(commandBuffer, 0, 1, &scissor);

                VkBuffer vertexBuffers[] = { vertexBuffer };
                VkDeviceSize offsets[] = { 0 };

                vkCmdBindVertexBuffers(commandBuffer, 0, 1, vertexBuffers, offsets);
                vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);
                vkCmdDrawIndexed(commandBuffer, (uint32_t)indices.size(), 1, 0, 0, 0);
            }
            vkCmdEndRenderPass(commandBuffer);
        }
        vkEndCommandBuffer(commandBuffer);

        app.submitAndPresent(waitNextImage, waitSubmission, waitFence, imageIndex);
        glfwPollEvents();
    }

    vkWaitForFences(app.getDevice(), 1, &waitFence, VK_TRUE, UINT64_MAX);
    vkDestroyBuffer(app.getDevice(), vertexBuffer, nullptr);
    vkDestroyBuffer(app.getDevice(), indexBuffer, nullptr);
    return 0;
}